/**
 * 
 */
package edu.arizona.bio5.onto.decomposition;

import java.util.Collection;
import java.util.Collections;
import java.util.HashSet;
import java.util.Iterator;
import java.util.Set;

import org.semanticweb.owlapi.apibinding.OWLManager;
import org.semanticweb.owlapi.model.OWLAxiom;
import org.semanticweb.owlapi.model.OWLEntity;
import org.semanticweb.owlapi.model.OWLOntology;
import org.semanticweb.owlapi.model.OWLOntologyCreationException;
import org.semanticweb.owlapi.model.OWLOntologyManager;
import org.semanticweb.owlapi.modularity.OntologySegmenter;

import uk.ac.manchester.cs.graph.Edge;
import uk.ac.manchester.cs.graph.alg.BFS;
import uk.ac.manchester.cs.graph.alg.TransitiveReductionAlgorithm;
import uk.ac.manchester.cs.graph.alg.TransitiveReductionAlgorithm.ReductionHandler;
import uk.ac.manchester.cs.graph.alg.VertexVisitor;
import uk.ac.manchester.cs.owlapi.modularity.ModuleType;
import uk.ac.manchester.cs.owlapi.modularity.SyntacticLocalityModuleExtractor;

import com.clarkparsia.owlapi.modularity.locality.LocalityClass;

import edu.arizona.bio5.onto.decomposition.locality.NonSyntLocalAxiomSignatureFinder;

/**
 * This algorithm keeps the entire ontology in memory but doesn't require updating
 * minimal seed signature-based labels
 * 
 * @author Pavel Klinov
 *
 */
public class IncrementalADAlgorithm {
	
	private final OWLOntologyManager mMgr;
	private final OWLOntology mOntology;
	/*
	 * Stores the number of affected atoms by the last change
	 */
	private int mAffAtomNumber = 0;
	
	public IncrementalADAlgorithm(OWLOntology ontology) {
		mMgr = OWLManager.createOWLOntologyManager();
		mOntology = ontology;
	}
	
	protected int getAffectedAtomNumber() {
		return mAffAtomNumber;
	}

	public AtomicDecomposition addAxiom(AtomicDecomposition ad, OWLAxiom axiom) {
		if (ad.getTermBasedIndex() == null) {
			ad.buildBasicTermIndex();
		}
		
		return processAdd(ad, axiom); 
	}

	private AtomicDecomposition processAdd(AtomicDecomposition ad, OWLAxiom axiom) {
		//First, see if any axioms are global
		Set<OWLAxiom> prevGlobalAx = findGlobalAxioms(ad.getModuleType());
		
		mOntology.getOWLOntologyManager().addAxiom(mOntology, axiom);
		
		Set<OWLAxiom> newGlobalAx = findGlobalAxioms(ad.getModuleType());
		
		if (!newGlobalAx.equals(prevGlobalAx)) {
			// FIXME What to do here?!
			ADAlgorithm decAlgo = new ChiaraDecompositionAlgorithm(ad.getModuleType());
			AtomicDecomposition newAd = decAlgo.decompose(mMgr, mOntology); 
			//build new term index
			newAd.buildBasicTermIndex();
			
			return newAd;
			
		} else {
			
			ad.setGlobalAxioms(newGlobalAx);
			
			if (!newGlobalAx.isEmpty()) {
				removeAffectedAtoms(ad, findAtomsForAxioms(ad, newGlobalAx, true, false));
			}
			// find the atoms whose PIs are affected by the change
			// find atoms whose signature contains some MGS of the new axiom
			// get all dependents of those atoms. These are the affected atoms
			Set<Atom> affAtoms = getAffectedAtoms(ad, axiom, newGlobalAx);
			
			mAffAtomNumber = affAtoms.size();
			// Remove that part from the existing AD
			removeAffectedAtoms(ad, affAtoms);
			// Re-decompose the affected part
			AtomicDecomposition newAd = redecompose(axiom, affAtoms, newGlobalAx, ad.getModuleType());
			// Merge the affected part back in. This requires ONLY adding
			// dependencies from the affected atoms to atoms in the unaffected part of the ontology.
			// Here MSS-labels of unaffected atoms could be used.
			merge(ad, newAd);
		}
		
		return ad;
	}
	
	private Set<OWLAxiom> findGlobalAxioms(ModuleType modType) {
		OntologySegmenter segmenter = new SyntacticLocalityModuleExtractor(mMgr, mOntology, modType);
		
		return segmenter.extract(Collections.<OWLEntity>emptySet());
	}

	private Set<Atom> findAtomsForAxioms(AtomicDecomposition ad, Set<OWLAxiom> axioms, boolean dependencies, boolean dependents) {
		Set<Atom> atoms = new HashSet<Atom>(axioms.size());
		
		for (OWLAxiom axiom : axioms) {
			Atom atom = ad.getAtomForAxiom(axiom);
			
			if (atom != null) {
				atoms.add(atom);
				
				if (dependencies) atoms.addAll(ad.getDependencies(atom));
				if (dependents) atoms.addAll(ad.getDependents(atom));
			}
		}
		
		return atoms;
	}

	private void merge(final AtomicDecomposition ad, final AtomicDecomposition newAd) {
		
		if (newAd == null) return;
		
		OntologySegmenter segmenter = new SyntacticLocalityModuleExtractor(mMgr, mOntology, ad.getModuleType());
		Set<OWLAxiom> newAxioms = new HashSet<OWLAxiom>(newAd.getGlobalAxioms());
		final Set<Atom> newAtoms = new HashSet<Atom>();
		
		for (Atom atom : newAd.getAtoms()) {
			newAxioms.addAll(atom.getAxioms());
		}
		//Copy new atoms and dependencies
		for (Atom topAtom : newAd.getTopAtoms()) {
			new BFS().search(newAd.getGraph(), topAtom, new VertexVisitor<Atom>(){

				@Override
				public boolean visit(Atom newAtom) {
					Atom newStart = null;
					Atom newEnd = null;
					
					if ((newStart = ad.getAtomForAxiom(newAtom.getAxioms().iterator().next())) == null) {
						newStart = ad.addAtom(newAtom.getAllAxioms());
					}
					else {//copy dependencies from unaffected atoms
						
					}
					
					newAtoms.add(newStart);
					
					for (Atom dep : newAd.getDependencies(newAtom, true)) {
						if ((newEnd = ad.getAtomForAxiom(dep.getAxioms().iterator().next())) == null) {
							newEnd = ad.addAtom(dep.getAllAxioms());
						}
						
						ad.addDependency(newStart, newEnd);
						newAtoms.add(newEnd);
					}
					
					return true;
				}});
		}
		
		//System.out.println("New atoms: " + newAtoms);
		// Add dependencies to unaffected atoms, if any
		for (Atom newAtom : newAtoms) {
			//Instead of extracting modules we could take advantage of MSS-labels
			Set<OWLAxiom> module = segmenter.extract(newAtom.getSignature());
			//See if the module includes any axioms which are not in the new AD
			for (OWLAxiom axiom : module) {
				if (!newAxioms.contains(axiom)) {
					//This either means a new dependency on a unaffected atom or that this axiom is the one being added 
					Atom unaffected = ad.getAtomForAxiom(axiom);
					
					if (unaffected != null) {
						ad.addDependency(newAtom, unaffected);
						//System.out.println(newAtom + " -> " + unaffected + "(" + unaffected.getAxioms() + ")");
					}
				}
			}
		}
		//Merge in the new term index
		//buildTermIndex(ad);
		updateTermIndex(ad, newAd);
		//Transitive reduction
		new TransitiveReductionAlgorithm<Atom, Object>(new ReductionHandler<Atom, Object>() {

			@Override
			public void onEdgeRemove(Edge<Atom, Object> edge) {
				ad.removeDependency(edge.getStart(), edge.getEnd());
			}})
			.compute(ad.getGraph(), newAtoms);
	}

	/*
	 * Updates the term index in the original decomposition by accounting for the new atoms
	 */
	private void updateTermIndex(AtomicDecomposition ad, AtomicDecomposition newAd) {
		
		for (Atom newAtom : newAd.getAtoms()) {
			for (OWLEntity term : newAtom.getSignature()) {
				Set<Atom> termAtoms = ad.getTermBasedIndex().get(term.getIRI().toURI());
				
				termAtoms = termAtoms == null ? new HashSet<Atom>() : termAtoms;
				termAtoms.add(ad.getAtomForAxiom(newAtom.getAxioms().iterator().next()));
				ad.getTermBasedIndex().put(term.getIRI().toURI(), termAtoms);
			}
		}		
	}

	private AtomicDecomposition redecompose(OWLAxiom addedAxiom, Set<Atom> affAtoms, Set<OWLAxiom> newGlobalAxioms, ModuleType modType) {
		//Re-decomposition should account for global axioms, of course
		Set<OWLAxiom> axioms = new HashSet<OWLAxiom>(newGlobalAxioms);
		
		for (Atom atom : affAtoms) {
			axioms.addAll(atom.getAxioms());
		}
		
		axioms.add(addedAxiom);
		
		try {
			//Re-decompose
			OWLOntology affectedOntology = mMgr.createOntology(axioms);
			ADAlgorithm decAlgo = new ChiaraDecompositionAlgorithm(modType);
			AtomicDecomposition ad = decAlgo.decompose(mMgr, affectedOntology); 
			//build new term index
			ad.buildBasicTermIndex();
			
			return ad;
			
		} catch (OWLOntologyCreationException e) {
			// Shouldn't happen
			throw new RuntimeException(e);
		}
	}

	/*
	 * TODO Save incoming edges from unaffected atoms, those need to be preserved
	 */
	private void removeAffectedAtoms(AtomicDecomposition ad, Set<Atom> affAtoms) {
		for (Atom atom : affAtoms) {
			ad.removeAtom(atom);
			//update the term index
			for (OWLEntity term : atom.getSignature()) {
				ad.getTermBasedIndex().get(term.getIRI().toURI()).remove(atom);
			}
		}
	}

	private Set<Atom> getAffectedAtoms(AtomicDecomposition ad, OWLAxiom axiom, Set<OWLAxiom> newGlobalAxioms) {
		NonSyntLocalAxiomSignatureFinder mgsFinder = new NonSyntLocalAxiomSignatureFinder(ad.getModuleType() == ModuleType.BOT ? LocalityClass.BOTTOM_BOTTOM : LocalityClass.TOP_TOP);
		Collection<Set<OWLEntity>> mgsSets = mgsFinder.findMinimalSeedSignatures(axiom, axiom.getSignature());
		Set<OWLEntity> globalSignature = new HashSet<OWLEntity>(newGlobalAxioms.size() * 3);

		for (OWLAxiom global : newGlobalAxioms) {
			globalSignature.addAll(global.getSignature());
		}
		
		//System.out.println("MGS: " + mgsSets);
		//System.out.println("Global signature: " + globalSignature);
		
		//Now get the immediately affected atoms
		Set<Atom> affAtoms = new HashSet<Atom>();
		
		for (Set<OWLEntity> mgs : mgsSets) {
			Set<Atom> atomsForMGS = new HashSet<Atom>();
			//Remove the signature of global axioms. We don't need it to pick atoms which are affected 
			mgs.removeAll(globalSignature);
			//find all atoms in which occur some term from the current MGS
			for (OWLEntity term : mgs) {
				Collection<Atom> atomsForTerm = ad.getTermBasedIndex().get(term.getIRI().toURI());
				
				if (atomsForTerm != null) {
					atomsForMGS.addAll(atomsForTerm);
				}
			}
			//ignore all atoms whose PI's signature does not contain the MGS
			for (Iterator<Atom> atomIter = atomsForMGS.iterator(); atomIter.hasNext();) {
				Atom candidateAff = atomIter.next();
				
				if (!candidateAff.getSignature().containsAll(mgs) && !ad.getPrincipalIdealSignature(candidateAff).containsAll(mgs)) {
					atomIter.remove();
				}
			}
			
			affAtoms.addAll(atomsForMGS);
		}
		//Now get all dependents and dependencies for each affected atom
		Set<Atom> allAffectedAtoms = new HashSet<Atom>();
		
		for (Atom affectedAtom : affAtoms) {
			if (ad.isBottomAtom(affectedAtom)) {
				allAffectedAtoms.addAll(ad.getDependents(affectedAtom));
			}
			else {
				allAffectedAtoms.addAll(ad.getRelatedAtoms(affectedAtom));
			}
		}
		
		//System.out.println("All affected atoms: " + allAffectedAtoms + ": ");
		//for (Atom aff : allAffectedAtoms) System.out.println(aff.getAxioms());
		
		return allAffectedAtoms;
	}
}